import asyncio
import time

from app.logger import logger
from app.agent import PlanningAgent, CoTAgent
from app.agent.manus import Manus
from app.agent.score import ScoreAgent
from app.flow.flow_factory import FlowFactory, FlowType
from app.web.chatbot.BaseChatBot import BaseChatBot


class FlowChatBot(BaseChatBot):
    __agents = {
            "planning": PlanningAgent(),
            "manus": Manus(),
            'cot': CoTAgent(),
            'score': ScoreAgent()
        }
    __flow = None

    def __init__(self):
        self.__flow = FlowFactory.create_flow(flow_type=FlowType.PLANNING, agents=self.__agents)

    async def ask(self, prompt):
        try:
            start_time = time.time()
            result = await asyncio.wait_for(
                self.__flow.execute(prompt),
                timeout=3600,  # 60 minute timeout for the entire execution
            )
            elapsed_time = time.time() - start_time
            logger.info(f"Request processed in {elapsed_time: .2f} seconds\n result: {result}")
            return result
        except asyncio.TimeoutError:
            logger.error("Request processing timed out after 1 hour")
            logger.info(
                "Operation terminated due to timeout. Please try a simpler request."
            )
        except KeyboardInterrupt:
            logger.info("Operation cancelled by user.")
        except Exception as e:
            logger.error(f"Error: {str(e)}")

